﻿
#include "alpc.h"

PFN_NtAlpcCreatePort ZwAlpcCreatePort = NULL;
PFN_NtAlpcConnectPort ZwAlpcConnectPort = NULL;
PFN_NtAlpcConnectPortEx ZwAlpcConnectPortEx = NULL;
PFN_NtAlpcAcceptConnectPort ZwAlpcAcceptConnectPort = NULL;
PFN_NtAlpcDisconnectPort ZwAlpcDisconnectPort = NULL;
PFN_NtAlpcSetInformation ZwAlpcSetInformation = NULL;
PFN_NtAlpcQueryInformation ZwAlpcQueryInformation = NULL;
PFN_NtAlpcSendWaitReceivePort ZwAlpcSendWaitReceivePort = NULL;
PFN_NtAlpcCreateSecurityContext ZwAlpcCreateSecurityContext = NULL;
PFN_NtAlpcCancelMessage ZwAlpcCancelMessage = NULL;
PFN_NtAlpcCreatePortSection ZwAlpcCreatePortSection = NULL;

PVOID LoadFunction(const wchar_t* strFunction) {
    UNICODE_STRING strName;
    RtlInitUnicodeString(&strName, strFunction);
    PVOID pFunction = MmGetSystemRoutineAddress(&strName);
    return pFunction;
}

NTSTATUS InitAlpcEnv()
{
#define X(Name) ((*(PVOID *)&Name = LoadFunction(L#Name)) == NULL)

    X(ZwAlpcCreatePort);
    X(ZwAlpcConnectPort);
    X(ZwAlpcConnectPortEx);
    X(ZwAlpcAcceptConnectPort);
    X(ZwAlpcDisconnectPort);
    X(ZwAlpcSetInformation);
    X(ZwAlpcQueryInformation);
    X(ZwAlpcSendWaitReceivePort);
    X(ZwAlpcCreateSecurityContext);
    X(ZwAlpcCancelMessage);
    X(ZwAlpcCreatePortSection);
    
    return STATUS_SUCCESS;
}

void UninitAlpcEnv()
{

}


VOID NTAPI AlpclInitUnicodeString(
    PUNICODE_STRING DestinationString,
    PCWSTR SourceString
)
{
    // 设置 UNICODE_STRING 的字段
    DestinationString->Buffer = (PWSTR)SourceString;
    DestinationString->Length = (USHORT)(wcslen(SourceString) * sizeof(WCHAR));  // 计算实际长度（不包括 NULL 字符）
    DestinationString->MaximumLength = (USHORT)((wcslen(SourceString) + 1) * sizeof(WCHAR)); // 计算最大长度（包括 NULL 字符）
}

HANDLE AlpcCreatePort(LPCWSTR strName)
{
    NTSTATUS Status = 0;
    HANDLE hPort = NULL;
    UNICODE_STRING strPortName = { 0 };
    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    ALPCPORT_ATTRIBUTES PortAttributes = { 0 };

    // 初始化 ALPC 端口属性
    PortAttributes.MaxMessageLength = ALPC_MAX_BUFLEN;

    if (NULL == strName || L'\0' == strName[0])
    {
        return NULL;
    }

    // 初始化端口名称
    AlpclInitUnicodeString(&strPortName, strName);

    // 设置对象属性
    InitializeObjectAttributes(
        &ObjectAttributes,
        &strPortName,
        0,
        NULL,
        NULL
    );

    Status = ZwAlpcCreatePort(&hPort, &ObjectAttributes, &PortAttributes);

    if (0 != Status) {
        if (hPort)
        {
            AlpcClose(hPort);
            hPort = NULL;
        }
        return NULL;
    }

    return hPort;
}


HANDLE AlpcConnectPort(LPCWSTR strName)
{
    // 连接到服务端
    NTSTATUS Status = 0;
    HANDLE hPort = NULL;
    // 端口名称（与服务端相同）
    UNICODE_STRING strPortName = { 0 };
    ALPCPORT_ATTRIBUTES PortAttributes = { 0 };
    ALPCPORT_MESSAGE ConnectMessage = { 0 };
    ConnectMessage.u1.s1.DataLength = 0;  // 无数据
    ConnectMessage.u1.s1.TotalLength = sizeof(ALPCPORT_MESSAGE);
    ConnectMessage.MessageId = 0;
    ULONG dwDataSize = sizeof(ALPCPORT_MESSAGE);

    PortAttributes.MaxMessageLength = ALPC_MAX_BUFLEN; // technically the hard limit for this is 65535, if no constrains you can use AlpcMaxAllowedMessageLength() to set this limit

    if (NULL == strName || L'\0' == strName[0])
    {
        return NULL;
    }

    // 初始化端口名称
    AlpclInitUnicodeString(&strPortName, strName);

    Status = ZwAlpcConnectPort(
        &hPort,                    // 返回的端口句柄
        &strPortName,           // 服务端端口名称
        NULL,                      // 客户端对象属性
        &PortAttributes,                      // 端口属性
        ALPC_SYNC_CONNECTION,      // 连接标志
        NULL,                      // 安全描述符
        &ConnectMessage, // 连接消息
        NULL,// 连接消息大小      
        NULL,                      // 输出消息属性
        NULL,                      // 输入消息属性
        NULL                   // 超时时间
    );

    if (0 != Status) {
        if (hPort)
        {
            AlpcClose(hPort);
            hPort = NULL;
        }
        return NULL;
    }

    return hPort;
}

HANDLE AlpcAcceptPort(HANDLE hServer, ULONG uMessageId)
{
    // 接受连接请求
    NTSTATUS Status = 0;
    HANDLE hClient = NULL;
    ALPCPORT_MESSAGE RecvData = { 0 };
    RecvData.u1.s1.DataLength = 0;  // 无数据
    RecvData.u1.s1.TotalLength = sizeof(ALPCPORT_MESSAGE);
    RecvData.MessageId = uMessageId;

    Status = ZwAlpcAcceptConnectPort(
        &hClient,
        hServer,
        ALPC_SYNC_CONNECTION,
        NULL,
        NULL,
        NULL,
        &RecvData,
        NULL,
        TRUE
    );

    if (0 != Status) {
        if (hClient)
        {
            AlpcClose(hClient);
            hClient = NULL;
        }
        return NULL;
    }

    return hClient;
}

VOID AlpcClose(HANDLE hPort)
{
    NTSTATUS Status = 0;
    Status = ZwAlpcDisconnectPort(hPort, 0);
    ZwClose(hPort);
}

NTSTATUS AlpcSendWaitReceive(HANDLE hPort, PALPCPORT_MESSAGE pSendData,
    PALPCPORT_MESSAGE pRecvData, ULONG* pdwRecvSize, ULONG dwTimeOut)
{
    NTSTATUS Status = 0;
    LARGE_INTEGER Timeout = { 0 };
    Timeout.QuadPart = -10000 * dwTimeOut;  // 1毫秒超时

    Status = ZwAlpcSendWaitReceivePort(
        hPort,                         // 端口句柄
        0,
        pSendData,                       // 发送消息
        NULL,                               // 发送消息的附加属性
        pRecvData,                    // 接收消息
        pdwRecvSize,                      // 接收缓冲区长度
        NULL,                               // 接收消息的附加属性
        -1 == dwTimeOut ? NULL : &Timeout                             // 超时设置
    );

    return Status;
}

NTSTATUS AlpcClientSendReceive(LPCWSTR strName, PALPCPORT_MESSAGE pSendData,
    PALPCPORT_MESSAGE pRecvData, ULONG* pdwRecvSize, ULONG dwTimeOut)
{
    NTSTATUS Status = FALSE;
    HANDLE hPort = AlpcConnectPort(strName);
    if (NULL == hPort)
    {
        return STATUS_UNSUCCESSFUL;
    }

    Status = AlpcSendWaitReceive(hPort, pSendData, pRecvData, pdwRecvSize, dwTimeOut);

    AlpcClose(hPort);

    return Status;
}